// ccmuxtest.cpp
// Copyright (C) 2008 Sam Bateman
// samuel.bateman@gmail.com
//
// This file is part of DVD On Fire.
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

#include <cerrno>
#include <cstdio>
#include <cstring>
#include <error.h>
#include <string>
#include <sys/stat.h>
#include <vector>
using namespace std;

typedef unsigned char byte;

// struct Data {{{
struct Data
{
    Data()
    : bin_fp(NULL), m2v_fp(NULL), out_fp(NULL),
      bytes_read(0), bytes_total(0),
      frames_read(0), frames_total(0)
    {}

    FILE         *bin_fp;
    FILE         *m2v_fp;
    FILE         *out_fp;
    vector<byte>  buffer;
    off_t         bytes_read;
    off_t         bytes_total;
    off_t         frames_read;
    off_t         frames_total;
};
// }}}

// findPacket(data, search) {{{
bool findPacket(Data *data, byte find)
{
    // 4-byte MPEG video codes
    byte search_code[] = { 0x00, 0x00, 0x01, find };
    byte frame_code [] = { 0x00, 0x00, 0x01, 0x00 };
    byte header_code[] = { 0xff, 0xff, 0xff, 0xff };

    // read bytes into a buffer
    data->bytes_read = 0;
    data->frames_read = 0;
    byte b;
    while (fread(&b, 1, 1, data->m2v_fp) == 1){
        data->buffer.push_back(b);
        ++data->bytes_read;
        ++data->bytes_total;
        header_code[0] = header_code[1];
        header_code[1] = header_code[2];
        header_code[2] = header_code[3];
        header_code[3] = b;

        // check if we have found a frame boundary or our search byte
        if (memcmp(header_code, frame_code, 4) == 0){
            ++data->frames_read;
            ++data->frames_total;
        }
        if (memcmp(header_code, search_code, 4) == 0){
            return true;
        }
    }

    // we haven't found it
    return false;
}
// }}}
// main(argc, argv) {{{
int main(int argc, char **argv)
{
    // check for command-line arguments
    if (argc < 3){
        printf("Usage: %s <bin file> <m2v file>\n", argv[0]);
        return 1;
    }

    // get the file size
    struct stat st;
    int r = stat(argv[2], &st);
    if (r != 0) error(1, errno, "Error with stat()");
    off_t size = st.st_size;

    // open the input and output files
    Data *data = new Data;
    data->bin_fp = fopen(argv[1], "rb");
    if (!data->bin_fp) error(1, errno, "Couldn't open file %s", argv[1]);
    data->m2v_fp = fopen(argv[2], "rb");
    if (!data->m2v_fp) error(1, errno, "Couldn't open file %s", argv[2]);
    data->out_fp = fopen("mtout.m2v", "wb");
    if (!data->out_fp) error(1, errno, "Couldn't open file mtout.m2v");

    // get timecode from gop header
    if (!findPacket(data, 0xb8)) error(1, 0, "Couldn't read timecode");
    byte time_code[4] = {0};
    fread(time_code, 1, 4, data->m2v_fp);
    int h = (int)time_code[0] / 0x04;
    int m = (int)time_code[1] / 0x10 + time_code[0] % 0x04 * 0x10;
    int s = (int)time_code[2] / 0x20 + time_code[1] % 0x08 * 0x08;
    int f = (int)time_code[3] / 0x80 + time_code[2] % 0x20 * 0x02;
    float fps = 30000 / 1001.0;
    int offset = h * 3600 * fps
               + m * 60   * fps
               + s        * fps
               + f;

    // account for frames before the gop header
    off_t gop_pos = data->bytes_total;
    data->bytes_total = 0;
    data->frames_total = 0;
    rewind(data->m2v_fp);
    while (true){
        if (!findPacket(data, 0x00)) error(1, 0, "Couldn't read gop frames");
        if (data->bytes_read >= gop_pos) break;
        --offset;
    }
    data->buffer.clear();
    data->bytes_total = 0;
    data->frames_total = 0;
    rewind(data->m2v_fp);
    if (offset > 0){
        for (int i = 0; i < offset; ++i){
            byte buf[2];
            fread(buf, 1, 2, data->bin_fp);
        }
    }

    // read and write mpeg frames
    bool read_cc = true;
    while (findPacket(data, 0xb8)){
        if (data->frames_read == 0) continue;
        if (offset < 0){
            offset += data->frames_read;
            continue;
        }

        // write the user data (captions)
        byte header[] = { 0x00, 0x00, 0x01, 0xb2,
                          0x43, 0x43, 0x01, 0xf8,
                          data->frames_read * 2 + 0x80 };
        fwrite(header, 1, 9, data->out_fp);
        for (int i = 0; i < data->frames_read; ++i){
            byte cc[] = { 0x00, 0x00 };
            if (!read_cc || fread(cc, 1, 2, data->bin_fp) != 2){
                read_cc = false;
                cc[0] = cc[1] = 0x80;
            }
            fprintf(data->out_fp, "\xff%c%c\xff%c%c",
                    cc[0], cc[1], 0x00, 0x00);
        }

        // write the mpeg data
        fwrite(&data->buffer[0], 1, data->buffer.size() - 4, data->out_fp);
        fflush(data->out_fp);
        data->buffer.clear();
        data->buffer.push_back(0x00);
        data->buffer.push_back(0x00);
        data->buffer.push_back(0x01);
        data->buffer.push_back(0xb8);

        // print progress
        printf("\rProgress: %.6f (GOP = %d)",
                (float)data->bytes_total / size,
                (int)data->frames_read);
        fflush(stdout);
    }

    // write the last chunk of mpeg data
    fwrite(&data->buffer[0], 1, data->buffer.size(), data->out_fp);
    data->buffer.clear();
    printf("\rProgress: 1.000000                \n");

    // clean up
    fclose(data->bin_fp);
    fclose(data->m2v_fp);
    fclose(data->out_fp);
    delete data;
}
// }}}
